import librosa
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import os
from torch.nn.utils.rnn import pad_sequence
import copy


def load_audio(file_path):
  y, sr = librosa.load(file_path, sr=22050)
  return y, sr

    
def get_all_waves_name(Wave_path):
  files = []
  for file in os.listdir(Wave_path):
      if os.path.isfile(os.path.join(Wave_path, file)):
          files.append(os.path.join(Wave_path, file))
  return files

def extract_mel_spectrogram(y, sr): #这个是用来提取梅尔频谱图的
    mel_spec = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=80)
    #log_mel_spec = librosa.power_to_db(mel_spec, ref=np.max)
    return mel_spec


class TTSDataset(Dataset): #数据加载器。
    def __init__(self, audio_files_path):
        self.audio_files = audio_files_path
    def __len__(self):
        return len(self.audio_files)
    def __getitem__(self, idx):
        audio_file = self.audio_files[idx]
        y, sr = load_audio(audio_file)
        mel_spec = extract_mel_spectrogram(y, sr)
        sample = {
            'audio_label': y,
            'sr': sr,
            'mel_spec': mel_spec,
        }
        return sample

def get_DataLoader(Wave_path = '../Data/Wave', batch_size=32, num_workers = 0):
    audio_files = get_all_waves_name(Wave_path) 
    dataset = TTSDataset(audio_files)
    #dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=collate_fn)
    dataloader = DataLoader(dataset, batch_size=32, shuffle=True, collate_fn = lambda x: x, pin_memory = False)#collate_fn=collate_fn)
    return dataloader


